import os
import numpy as np
import cv2
import pylab
import Params

def flipUd(arr):
    ''' vertically invert given matrix '''
    arr2 = np.zeros_like(arr, dtype='float32')
    arr2[:, :] = arr[::-1, :]
    return arr2

def side_by_side(arr1, arr2):
    ''' compile two images side by side '''
    shape1 = np.shape(arr1)
    shape2 = np.shape(arr2)

    if len(shape1) == 2:
        bk = np.array(arr1)
        arr1 = np.zeros((shape1[0], shape1[1], 3), dtype=bk.dtype)
        for i in range(3): arr1[:,:,i] = bk
        shape1 = arr1.shape
    if len(shape2) == 2:
        bk = np.array(arr2)
        arr2 = np.zeros((shape2[0], shape2[1], 3), dtype=bk.dtype)
        for i in range(3): arr2[:,:,i] = bk
        shape2 = arr2.shape

    if shape1[2] != shape2[2]:
        raise Exception("Color dimensions must match")
    elif arr1.dtype != arr2.dtype:
        print arr1.dtype, arr2.dtype
        raise Exception("Data types must match")
    elif shape1[0] != shape2[0]:
        raise Exception("Height must match")

    output = np.zeros((shape1[0], shape1[1] + shape2[1], shape1[2]), dtype=arr1.dtype)
    output[:, :shape1[1], :] = arr1
    output[:, shape1[1]:, :] = arr2
    return output

def writeImage(fn, array, limit=False, limitsize=0, flip_img=False):
    ''' write image from array to fn after fix '''
    
    # create needed folders
    try:
        os.makedirs(fn[:(len(fn) - fn[::-1].find('/'))])
    except OSError:
        pass;
    
    shape = np.shape(array)
    if len(shape) == 2:
        array = cv2.cvtColor(np.array(array), cv2.cv.CV_GRAY2BGR)

    if np.max(array) < 255:
        array = 255 * array

    if array.dtype != 'uint8':
        array = array.astype('uint8')

    
    yfail = xfail = False
    if limit and limitsize != 0:
        # print 'trying to limit'
        # print 'before:', array.shape
        while not (yfail and xfail):
            # print yfail, xfail
            
            # try cut y
            if np.sum(array[-limitsize:, :] != 127) == 0:
                array = array[:-limitsize, :]
                yfail = False
            elif np.sum(array[:limitsize, :] != 127) == 0:
                array = array[limitsize:, :]
                yfail = False
            else: yfail = True
            
            # try cut x
            if np.sum(array[:, -limitsize:] != 127) == 0:
                array = array[:, :-limitsize]
                xfail = False
            elif np.sum(array[:, :limitsize] != 127) == 0:
                array = array[:, limitsize:]
                xfail = False
            else: xfail = True
    
    # print 'after:', array.shape
    
    if flip_img:
        array = flipUd(array)
    cv2.imwrite(fn, array)
    print 'write', fn
    print np.sum(array)

def drawable(array, inv=False):
    shape = np.shape(array)
    if array.dtype != 'uint8':
        array = array.astype('uint8')

    if len(shape) == 2:
        array = cv2.cvtColor(np.array(array), cv2.cv.CV_GRAY2BGR)

    if inv:
        array = 255 * np.ones_like(array) - array
    
    return array

def pyrUp(arr):
    """ Increase size of array without smoothing """
    sh = np.shape(arr);
    narr = np.zeros((sh[0] * 2, sh[1] * 2), dtype=arr.dtype);
    narr[::2, ::2] = arr
    narr[1::2, ::2] = arr
    narr[::2, 1::2] = arr
    narr[1::2, 1::2] = arr
    return narr;
   
def top_and_bottom(arr1, arr2):
    ''' compile two images side by side '''
    shape1 = np.shape(arr1)
    shape2 = np.shape(arr2)

    if arr1.dtype != 'uint8':
        arr1 = arr1.astype('uint8')
    if arr2.dtype != 'uint8':
        arr2 = arr2.astype('uint8')
    
    if len(shape1) == 2:
        arr1 = cv2.cvtColor(np.array(arr1), cv2.cv.CV_GRAY2BGR)
        shape1 = np.shape(arr1)
    if len(shape2) == 2:
        arr2 = cv2.cvtColor(np.array(arr2), cv2.cv.CV_GRAY2BGR)
        shape2 = np.shape(arr2)
        
    if shape1[2] != shape2[2]:
        raise Exception("Color dimensions must match")
    elif arr1.dtype != arr2.dtype:
        raise Exception("Data types must match")
    elif shape1[1] != shape2[1]:
        raise Exception("Width must match")

    output = np.zeros((shape1[0] + shape2[0], shape1[1], shape1[2]), dtype=arr1.dtype)
    output[:shape1[0], :, :] = arr1
    output[shape1[0]:, :, :] = arr2
    return output

def remove_zero(arr):
    return arr[np.where(arr != 0)]

def resize(original, destination, method='zoom'):
    current = np.shape(original)
    colors = current[2]
    current = np.array(current)[0:2]
    
    if method == 'stretch':
        dest_size = destination
        return cv2.resize(original, dest_size)
    
    elif method == 'zoom':
        dest_ratio = (float(destination[0]) / current[0], float(destination[1]) / current[1])
        
        limit = np.argmin(dest_ratio)
        dest_size = np.array(current) * dest_ratio[limit]
        dest_size = (int(round(dest_size[1])), int(round(dest_size[0])))

        dest = cv2.resize(original, dest_size)
        l = len(destination)
        if l == 2:
            destination = (destination[0], destination[1], colors)
        output = np.zeros(destination, dtype=original.dtype) 

        if limit == 0:
            # height is same, width different
            dest_w = dest.shape[1]
            out_w = output.shape[1]
            output[:, (out_w - dest_w) / 2:(out_w - dest_w) / 2 + dest_w] = dest
            return output
        else:
            # width is same, height different
            dest_h = dest.shape[0]
            out_h = output.shape[0]
            output[(out_h - dest_h) / 2:(out_h - dest_h) / 2 + dest_h, :] = dest
            return output

    return None

base_disp = "%s/output/disp_%s_%s_%05d.png"
def fetch_disp(alg, frame, lib='tsukuba'):
    ''' fetch a single image from disk with a particular base '''

    factor = 1.
    if alg != "var" and alg != "sad":
        max_disp = 16*6
        factor = (255.0/(max_disp * 16)) * 16;
    elif alg == 'var':
        max_disp = 16*6.
        factor = 256/max_disp;
        

    img = None
    if lib == 'tsukuba':
        img = cv2.imread(base_disp % (lib, lib, alg, frame))
    else:
        img = cv2.imread(base_disp % (lib, lib, alg, frame))

    if img is None:
        raise Exception("Could not open image\n" + base_disp % (lib, lib, alg, frame))
    img = (img / factor).astype('uint8')

    return img

base_orig_tsukuba = "tsukuba/%s/tsukuba_daylight_%s_%05d.png"
base_orig_kitti = "kitti/%s/%s_%010d.png"
def fetch_orig(frame, lib='tsukuba', side='L'):
    ''' fetch a single image from disk with a particular base '''
    full = ('left' if side == 'L' else 'right')
    if lib == 'tsukuba':
        img = cv2.imread(base_orig_tsukuba % (full, side, frame))
    else:
        img = cv2.imread(base_orig_kitti % (full, side, frame))
    return img


base_ground_tsukuba = "tsukuba/ground truth/tsukuba_disparity_L_%05d.png"
# base_ground_kitti = "kitti/laser/disp_kitti_la_%05d.png"
base_ground_kitti = "kitti/laser/disp_kitti_la_%05d.npz"

def fetch_ground(frame, lib='tsukuba'):
    ''' fetch a single image from disk with a particular base '''
    if lib == 'tsukuba':
        img = cv2.imread(base_ground_tsukuba % frame)
    	if len(img.shape) > 2: img = img[:,:,0]
    else:
        data = np.load(base_ground_kitti % frame)
        img = data['ground']
        # img = cv2.imread(base_ground_kitti % frame)
    return img

def fetch_simulate(filename):
    ''' read simulation image '''
    fn = '%s/%s' % (Params.base_folder, filename)
    img = cv2.imread(fn)
    if len(img.shape) > 2: img = img[:,:,0]
    # img = cv2.cvtColor(img, cv2.cv.CV_BGR2GRAY)
    out = np.zeros_like(img, dtype='uint8')
    out[np.where(img > 1)] = 255
    return out

def fetch_simulate_gps(filename):
    ''' read simulation image '''
    print 'filename', filename
    fn = '%s/%s' % (Params.base_folder, filename)
    img = cv2.imread(fn)
    if len(img.shape) > 1: img = img[:, :, 0]
    # img = cv2.cvtColor(img, cv2.cv.CV_BGR2GRAY)
    out = np.zeros_like(img, dtype='uint8')
    out[np.where(img > 0)] = 255
    return out

def save_latex(x, y, labels, filename, xlabel=None, ylabel=None, width=1.0, draw=False, legend_title=None, xlim=None, ylim=[0,1], Z=None, double=False, article=False, yticks=None, xticks=None, bar=False, model=None, yfactor=1.0):
    pts = 245.71811 if article else 452.9679
    folder = Params.article_folder if article else Params.thesis_folder

    fig_width_pt = width * pts              # Get this from LaTeX using \showthe\columnwidth
    inches_per_pt = 1.0/72.27               # Convert pt to inch
    golden_mean = (pylab.sqrt(5)-1.0)/2.0   # Aesthetic ratio
    fig_width = fig_width_pt*inches_per_pt  # width in inches
    if double: fig_width *= 2
    fig_height = fig_width*golden_mean*yfactor      # height in inches
    fig_size =  [fig_width,fig_height]
    params = {'backend': 'ps',
              'axes.labelsize': 9,
              'text.fontsize': 8,
              'axes.titlesize': 9,
              'legend.fontsize': 7,
              'xtick.labelsize': 8,
              'ytick.labelsize': 8,
#               'text.usetex': True,
#               'text.latex.unicode': True,
              'svg.fonttype': 'none',
              'figure.figsize': fig_size}
    pylab.rcParams.update(params)

    # Plot data
    pylab.figure(1)
    pylab.clf()
    xoffset = .15
    yoffset = .2
    pylab.figure(figsize=(fig_width,fig_height))
    pylab.axes([xoffset,yoffset,0.95-xoffset,0.95-yoffset])

    if model is not None:
        pylab.bar(model.bins[:-1], model.count, color='r', linewidth=0, width=model.bins[1]-model.bins[0])
        pylab.hold(True)
        
    if y is not None:
        if bar:
            pylab.bar(x, y, width=(x[1]-x[0]))
        else:
            pylab.plot(x, y)
    
    pylab.grid()
    pylab.xlabel(xlabel)
    pylab.ylabel(ylabel)
    
    if Z is not None: 
        pylab.hold(True)
        pylab.plot([Z, Z], [0, 1], c='k', ls='dashed')
    
    if xticks is not None:
        pylab.xticks(xticks)
    if xlim is not None:
        pylab.xlim(xlim)
    if ylim is not None:
        pylab.ylim(ylim)
        if yticks is not None:
            pylab.yticks(yticks)
        elif ylim == [0, 1]:
            pylab.yticks([0, 0.5, 1])
    # pylab.legend(labels, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)
    if labels is not None:
        pylab.legend(labels, loc='upper right', title=legend_title)
    print 'saving picture', '%s/%s' % (folder,filename)
    pylab.savefig('%s/%s.svg' % (folder, filename), transparent=True)
    if draw:
        pylab.show()
